import akshare as ak
import pandas as pd
from pyecharts import options as opts
from pyecharts.charts import Kline

def get_shanghai_index_data(start_date):
    # 获取上证指数日线数据
    stock_zh_index_daily_df = ak.stock_zh_index_daily(symbol="sh000001")
    # 将 date 列转换为日期时间类型
    stock_zh_index_daily_df['date'] = pd.to_datetime(stock_zh_index_daily_df['date'])
    # 根据起始日期筛选数据
    recent_data = stock_zh_index_daily_df[stock_zh_index_daily_df['date'] >= start_date]
    return recent_data

def draw_kline_chart(data):
    # 准备 K 线图所需的数据
    kline_data = data[['open', 'close', 'low', 'high']].values.tolist()
    # 从 date 列获取日期数据并转换为字符串列表
    x_data = data['date'].dt.strftime('%Y-%m-%d').tolist()

    # 创建 K 线图对象
    kline = (
        Kline()
        .add_xaxis(x_data)
        .add_yaxis("上证指数", kline_data)
        .set_global_opts(
            title_opts=opts.TitleOpts(title="上证指数 K 线图"),
            xaxis_opts=opts.AxisOpts(type_="category"),
            yaxis_opts=opts.AxisOpts(
                splitarea_opts=opts.SplitAreaOpts(
                    is_show=True, areastyle_opts=opts.AreaStyleOpts(opacity=1)
                ),
            ),
            datazoom_opts=[opts.DataZoomOpts()],
            toolbox_opts=opts.ToolboxOpts(is_show=True)
        )
    )
    return kline

if __name__ == "__main__":
    # 自定义起始日期
    start_date = "2024-01-01"
    # 获取数据
    shanghai_data = get_shanghai_index_data(start_date)
    # 绘制 K 线图
    kline_chart = draw_kline_chart(shanghai_data)
    # 渲染为 HTML 文件
    kline_chart.render("shanghai_index_kline.html")